#include "BuildDatabase.h"

#include "BaseCommand.h"
#include "FeaturesFilter.h"
#include "Paths.h"
#include "building/CompileCommand.h"
#include "exceptions/CompilationDatabaseException.h"
#include "utils/DynamicLibraryUtils.h"
#include "utils/JsonUtils.h"
#include "utils/StringUtils.h"
#include "utils/GrpcUtils.h"
#include "utils/GenerationUtils.h"

#include "loguru.h"

#include <functional>
#include <queue>
#include <unordered_map>
#include <utility>

const std::string BuildDatabase::BITS_32_FLAG = "-m32";

BuildDatabase::BuildDatabase(
        fs::path serverBuildDir,
        fs::path buildCommandsJsonPath,
        fs::path linkCommandsJsonPath,
        fs::path compileCommandsJsonPath,
        utbot::ProjectContext projectContext
) : serverBuildDir(std::move(serverBuildDir)),
    buildCommandsJsonPath(std::move(buildCommandsJsonPath)),
    linkCommandsJsonPath(std::move(linkCommandsJsonPath)),
    compileCommandsJsonPath(std::move(compileCommandsJsonPath)),
    projectContext(std::move(projectContext)) {
}

BuildDatabase::BuildDatabase(BuildDatabase *baseBuildDatabase) :
        serverBuildDir(baseBuildDatabase->serverBuildDir),
        projectContext(baseBuildDatabase->projectContext),
        buildCommandsJsonPath(baseBuildDatabase->buildCommandsJsonPath),
        linkCommandsJsonPath(baseBuildDatabase->linkCommandsJsonPath),
        compileCommandsJsonPath(baseBuildDatabase->compileCommandsJsonPath) {
}

fs::path BuildDatabase::createExplicitObjectFileCompilationCommand(const std::shared_ptr<ObjectFileInfo> &objectInfo) {
    if (Paths::isSourceFile(objectInfo->getSourcePath())) {
        auto outputFile = objectInfo->getOutputFile();
        auto tmpObjectFileName =
                Paths::createTemporaryObjectFile(outputFile, objectInfo->getSourcePath());
        objectInfo->setOutputFile(tmpObjectFileName);
        // redirect existing compilation command to temporary file
        LOG_IF_S(ERROR, CollectionUtils::containsKey(objectFileInfos, tmpObjectFileName))
        << "Temporary object file name generated by UTBot is already present in the "
           "project: "
        << tmpObjectFileName;
        objectInfo->linkUnit = outputFile;
        objectFileInfos[tmpObjectFileName] = objectInfo;
        return tmpObjectFileName;
    } else {
        return objectInfo->getSourcePath();
    }
}

void BuildDatabase::createClangCompileCommandsJson() {
    CollectionUtils::MapFileTo<nlohmann::json> fileCompileCommands;
    for (const auto &[sourcePath, objectInfos]: sourceFileInfos) {
        const std::shared_ptr<ObjectFileInfo> &objectInfo = objectInfos.front();
        fileCompileCommands[sourcePath] = {{"directory", objectInfo->command.getDirectory()},
                                           {"command",   objectInfo->command.toString()},
                                           {"file",      objectInfo->command.getSourcePath()}};
    }

    nlohmann::json compileCommandsSingleFilesJson;
    for (const auto &compileCommand: fileCompileCommands) {
        compileCommandsSingleFilesJson.push_back(compileCommand.second);
    }

    fs::path clangCompileCommandsJsonPath = CompilationUtils::getClangCompileCommandsJsonPath(buildCommandsJsonPath);
    JsonUtils::writeJsonToFile(clangCompileCommandsJsonPath, compileCommandsSingleFilesJson);
    compilationDatabase = CompilationUtils::getCompilationDatabase(clangCompileCommandsJsonPath);
}

void BuildDatabase::mergeLibraryOptions(std::vector<std::string> &jsonArguments) const {
    for (auto it = jsonArguments.begin(); it != jsonArguments.end(); it++) {
        if (*it == DynamicLibraryUtils::libraryDirOption || *it == DynamicLibraryUtils::linkFlag) {
            auto next = std::next(it);
            *it += *next;
            *next = "";
        }
    }
    CollectionUtils::erase(jsonArguments, "");
}

namespace {
    CollectionUtils::OrderedFileSet collectLibraryDirs(const utbot::BaseCommand &command) {
        using namespace DynamicLibraryUtils;
        CollectionUtils::OrderedFileSet libraryDirs;
        for (std::string const &argument: command.getCommandLine()) {
            auto optionalLibraryPath = getLibraryAbsolutePath(argument, command.getDirectory());
            if (optionalLibraryPath.has_value()) {
                libraryDirs.insert(optionalLibraryPath.value());
            }
            if (StringUtils::startsWith(argument, libraryDirOptionWl)) {
                auto commaSeparated = StringUtils::split(argument, ',');
                bool isRpathNext = false;
                for (auto part: commaSeparated) {
                    if (part == rpathFlag) {
                        isRpathNext = true;
                        continue;
                    }
                    if (isRpathNext) {
                        isRpathNext = false;
                        libraryDirs.insert(part);
                    }
                }
            }
        }
        return libraryDirs;
    }

    CollectionUtils::MapFileTo<std::string> collectLibraryNames(const utbot::BaseCommand &command) {
        using namespace DynamicLibraryUtils;

        CollectionUtils::MapFileTo<std::string> libraryNames;

        for (const auto &argument: command.getCommandLine()) {
            if (Paths::isSharedLibraryFile(argument) && argument != command.getOutput() &&
                !StringUtils::startsWith(argument, libraryDirOptionWl)) {
                libraryNames.emplace(argument, argument);
            }
            if (StringUtils::startsWith(argument, linkFlag)) {
                std::string libraryName = argument.substr(linkFlag.length());
                std::string archiveFile = "lib" + libraryName + ".a";
                std::string sharedObjectFile = "lib" + libraryName + ".so";
                libraryNames.emplace(sharedObjectFile, argument);
                libraryNames.emplace(archiveFile, argument);
            }
        }
        return libraryNames;
    }
}

void BuildDatabase::addLibrariesForCommand(utbot::BaseCommand &command,
                                           BaseFileInfo &info,
                                           sharedLibrariesMap &sharedLibraryFiles,
                                           bool objectFiles) {
    if (command.isArchiveCommand()) {
        return;
    }
    auto libraryDirs = collectLibraryDirs(command);
    auto libraryNames = collectLibraryNames(command);
    std::unordered_map<std::string, fs::path> argumentToFile;
    for (auto const &[libraryName, argument]: libraryNames) {
        fs::path name = libraryName;
        for (auto const &libraryDir: libraryDirs) {
            if (CollectionUtils::containsKey(sharedLibraryFiles, name)) {
                if (CollectionUtils::containsKey(sharedLibraryFiles.at(name), libraryDir)) {
                    name = sharedLibraryFiles.at(name).at(libraryDir);
                }
            }
            fs::path fullPath = Paths::getFileFullPath(name, libraryDir);
            if (CollectionUtils::containsKey(targetInfos, fullPath)) {
                info.addFile(fullPath);
                LOG_IF_S(WARNING, objectFiles) << "Object file " << command.getOutput()
                                               << " has library dependencies: " << fullPath;
                argumentToFile[argument] = fullPath;
            } else {
                info.installedFiles.insert(fullPath);
            }
        }
    }
    for (auto &argument: command.getCommandLine()) {
        if (CollectionUtils::containsKey(argumentToFile, argument)) {
            argument = argumentToFile[argument];
        }
    }
}

const fs::path &BuildDatabase::getCompileCommandsJson() {
    return compileCommandsJsonPath;
}

const fs::path &BuildDatabase::getLinkCommandsJson() {
    return linkCommandsJsonPath;
}

std::vector<std::shared_ptr<BuildDatabase::ObjectFileInfo>>
BuildDatabase::getAllCompileCommands() const {
    std::vector<std::shared_ptr<ObjectFileInfo>> result;
    for (auto &[file, compilationUnit]: objectFileInfos) {
        result.emplace_back(compilationUnit);
    }
    return result;
}

fs::path BuildDatabase::getObjectFile(const fs::path &sourceFile) const {
    if (!CollectionUtils::containsKey(sourceFileInfos, sourceFile)) {
        std::string message = "Couldn't find object file for current source file " +
                              sourceFile.string();
        LOG_S(ERROR) << message;
        throw CompilationDatabaseException(message);
    }
    auto objectInfo = sourceFileInfos.at(sourceFile)[0];
    return objectInfo->getOutputFile();
}

CollectionUtils::FileSet BuildDatabase::getArchiveObjectFiles(const fs::path &archive) const {
    if (Paths::isGtest(archive)) {
        return {};
    }
    if (!CollectionUtils::containsKey(targetInfos, archive)) {
        std::string message = "Couldn't find current archive file linkage information for " + archive.string();
        LOG_S(ERROR) << message;
        throw CompilationDatabaseException(message);
    }
    std::shared_ptr<TargetInfo> targetInfo = targetInfos.at(archive);
    CollectionUtils::FileSet result;
    for (const auto &file: targetInfo->files) {
        if (Paths::isLibraryFile(file)) {
            auto archiveObjectFiles = getArchiveObjectFiles(file);
            CollectionUtils::extend(result, archiveObjectFiles);
        } else {
            const fs::path &sourcePath = getClientCompilationObjectInfo(file)->getSourcePath();
            if (Paths::isSourceFile(sourcePath)) {
                result.insert(file);
            } else {
                LOG_S(WARNING) << "Skipping not c/c++ source file: " << sourcePath.string();
            }
        }
    }
    return result;
}

CollectionUtils::FileSet BuildDatabase::getArchiveTargetFiles(const fs::path &archive) const {
    if (Paths::isGtest(archive)) {
        return {};
    }
    if (!CollectionUtils::containsKey(targetInfos, archive)) {
        std::string message = "Couldn't find current archive file linkage information for " + archive.string();
        LOG_S(ERROR) << message;
        throw CompilationDatabaseException(message);
    }
    std::shared_ptr<TargetInfo> targetInfo = targetInfos.at(archive);
    CollectionUtils::FileSet result = {archive};
    for (const auto &file: targetInfo->files) {
        if (Paths::isLibraryFile(file)) {
            auto archiveObjectFiles = getArchiveTargetFiles(file);
            result.insert(file);
            CollectionUtils::extend(result, archiveObjectFiles);
        }
    }
    return result;
}

fs::path BuildDatabase::getRootForSource(const fs::path &path) const {
    fs::path normalizedPath = Paths::normalizedTrimmed(path);
    if (Paths::isSourceFile(normalizedPath)) {
        if (!CollectionUtils::containsKey(sourceFileInfos, normalizedPath)) {
            std::string message =
                    "No source file found in compile_commands.json: " + path.string();
            LOG_S(ERROR) << message;
            throw CompilationDatabaseException(message);
        }
        auto const &sourceFileInfo = sourceFileInfos.at(normalizedPath);

        auto linkUnit = sourceFileInfo[0]->linkUnit;
        if (linkUnit.empty()) {
            std::string message =
                    "No executable or library found for current source file in link_commands.json: " + path.string();
            LOG_S(ERROR) << message;
            throw CompilationDatabaseException(message);
        }
        return getRootForSource(linkUnit);
    } else {
        if (!CollectionUtils::containsKey(targetInfos, normalizedPath)) {
            std::string message =
                    "No executable or library found in link_commands.json: " + path.string();
            LOG_S(ERROR) << message;
            throw CompilationDatabaseException(message);
        }
        auto linkUnit = targetInfos.at(normalizedPath);
        if (!linkUnit->parentLinkUnits.empty()) {
            return getRootForSource(linkUnit->parentLinkUnits[0]);
        } else {
            return linkUnit->getOutput();
        }
    }
}

fs::path BuildDatabase::getRootForFirstSource() const {
    if (sourceFileInfos.empty()) {
        std::string message = "Source files not found";
        LOG_S(ERROR) << message;
        throw CompilationDatabaseException(message);
    } else {
        return getRootForSource(sourceFileInfos.begin()->first);
    }
}

fs::path BuildDatabase::getBitcodeForSource(const fs::path &sourceFile) const {
    fs::path serverBuildObjectFilePath = newDirForFile(sourceFile);
    return Paths::addExtension(serverBuildObjectFilePath, ".bc");
}

fs::path BuildDatabase::getBitcodeFile(const fs::path &filepath) const {
    auto objectInfo = sourceFileInfos.find(filepath);
    if (objectInfo != sourceFileInfos.end()) {
        return getBitcodeForSource(objectInfo->second[0]->getSourcePath());
    } else {
        auto objectInfo = objectFileInfos.find(filepath);
        if (objectInfo != objectFileInfos.end()) {
            return getBitcodeForSource(objectInfo->second->getSourcePath());
        } else {
            auto targetInfo = targetInfos.find(filepath);
            if (targetInfo != targetInfos.end()) {
                fs::path movedFile = newDirForFile(filepath);
                return getCorrespondingBitcodeFile(movedFile);
            }
            return getCorrespondingBitcodeFile(filepath);
        }
    }
}


std::shared_ptr<BuildDatabase::ObjectFileInfo>
BuildDatabase::getClientCompilationObjectInfo(const fs::path &filepath) const {
    if (!CollectionUtils::contains(objectFileInfos, filepath)) {
        std::string message = "Object file not found in compilation_commands.json: " + filepath.string();
        LOG_S(ERROR) << message;
        throw CompilationDatabaseException(message);
    }
    return objectFileInfos.at(filepath);
}

std::shared_ptr<BuildDatabase::ObjectFileInfo>
BuildDatabase::getClientCompilationSourceInfo(const fs::path &filepath) const {
    if (!CollectionUtils::contains(sourceFileInfos, filepath)) {
        std::string message = "Source file not found in compilation_commands.json: " + filepath.string();
        LOG_S(ERROR) << message;
        throw CompilationDatabaseException(message);
    }
    LOG_IF_S(DEBUG, sourceFileInfos.at(filepath).size() > 1) << "More than one compile command for: " << filepath;
    return sourceFileInfos.at(filepath)[0];
}

std::shared_ptr<const BuildDatabase::ObjectFileInfo>
BuildDatabase::getClientCompilationUnitInfo(const fs::path &filepath) const {
    if (Paths::isSourceFile(filepath)) {
        return getClientCompilationSourceInfo(filepath);
    }
    if (Paths::isObjectFile(filepath)) {
        return getClientCompilationObjectInfo(filepath);
    }
    std::string message = "File is not a compilation unit or an object file: " + filepath.string();
    LOG_S(ERROR) << message;
    throw CompilationDatabaseException(message);
}


[[nodiscard]] bool BuildDatabase::hasUnitInfo(const fs::path &filepath) const {
    return CollectionUtils::contains(sourceFileInfos, filepath) || CollectionUtils::contains(objectFileInfos, filepath);
}

std::shared_ptr<const BuildDatabase::TargetInfo> BuildDatabase::getClientLinkUnitInfo(const fs::path &filepath) const {
    if (Paths::isSourceFile(filepath)) {
        auto compilationInfo = getClientCompilationUnitInfo(filepath);
        return targetInfos.at(compilationInfo->linkUnit);
    }
    if (CollectionUtils::containsKey(targetInfos, filepath)) {
        return targetInfos.at(filepath);
    }
    std::string message = "File is not in link_commands.json: " +
                          filepath.string();
    LOG_S(ERROR) << message;
    throw CompilationDatabaseException(message);
}

static inline bool maybe64bits(const fs::path &path) {
    return StringUtils::contains(path.string(), "64");
}

bool BuildDatabase::ObjectFileInfo::conflictPriorityMore(
        const std::shared_ptr<BuildDatabase::ObjectFileInfo> &left,
        const std::shared_ptr<BuildDatabase::ObjectFileInfo> &right) {
    bool leftContains64 = maybe64bits(left->getOutputFile());
    bool rightContains64 = maybe64bits(right->getOutputFile());
    if (leftContains64 == rightContains64) {
        return left->getOutputFile().string() < right->getOutputFile().string();
    }
    return leftContains64;
}

fs::path BuildDatabase::getCorrespondingBitcodeFile(const fs::path &filepath) {
    return Paths::replaceExtension(filepath, ".bc");
}

bool BuildDatabase::isFirstObjectFileForSource(const fs::path &objectFilePath) const {
    fs::path sourceFile = getClientCompilationUnitInfo(objectFilePath)->getSourcePath();
    fs::path firstObjectFileForSource = getClientCompilationUnitInfo(sourceFile)->getOutputFile();
    return objectFilePath == firstObjectFileForSource;
}

BuildDatabase::KleeFilesInfo::KleeFilesInfo(fs::path kleeFile) : kleeFile(std::move(kleeFile)) {
    fs::create_directories(this->kleeFile.parent_path());
}

void BuildDatabase::KleeFilesInfo::setCorrectMethods(std::unordered_set<std::string> correctMethods) {
    this->correctMethods = std::move(correctMethods);
}

bool BuildDatabase::KleeFilesInfo::isCorrectMethod(const std::string &method) {
    if (allAreCorrect) {
        return true;
    }
    return CollectionUtils::contains(correctMethods, method);
}

fs::path BuildDatabase::KleeFilesInfo::getKleeFile() {
    return getKleeFile("");
}

fs::path BuildDatabase::KleeFilesInfo::getKleeBitcodeFile() {
    return getKleeBitcodeFile("");
}

fs::path BuildDatabase::KleeFilesInfo::getKleeFile(const std::string &methodName) {
    return kleeFile;
}

fs::path BuildDatabase::KleeFilesInfo::getKleeBitcodeFile(const std::string &methodName) {
    return getCorrespondingBitcodeFile(getKleeFile());
}

void BuildDatabase::KleeFilesInfo::setAllAreCorrect(bool allAreCorrect) {
    this->allAreCorrect = allAreCorrect;
}

fs::path const &BuildDatabase::ObjectFileInfo::getDirectory() const {
    return command.getDirectory();
}

fs::path BuildDatabase::ObjectFileInfo::getSourcePath() const {
    return command.getSourcePath();
}

fs::path BuildDatabase::ObjectFileInfo::getOutputFile() const {
    return command.getOutput();
}

void BuildDatabase::ObjectFileInfo::setOutputFile(const fs::path &file) {
    command.setOutput(file);
}

void BuildDatabase::BaseFileInfo::addFile(fs::path file) {
    files.insert(std::move(file));
}

bool BuildDatabase::ObjectFileInfo::is32bits() const {
    return CollectionUtils::contains(command.getCommandLine(), BITS_32_FLAG);
}

fs::path BuildDatabase::TargetInfo::getOutput() const {
    if (commands.empty()) {
        std::string message = "There are no targets";
        LOG_S(ERROR) << message;
        throw CompilationDatabaseException(message);
    }
    return commands[0].getOutput();
}

std::vector<std::shared_ptr<BuildDatabase::TargetInfo>> BuildDatabase::getRootTargets() const {
    return CollectionUtils::filterOut(
            CollectionUtils::getValues(targetInfos),
            [](const std::shared_ptr<const BuildDatabase::TargetInfo> &linkUnitInfo) {
                return !linkUnitInfo->parentLinkUnits.empty();
            });
}

std::vector<std::shared_ptr<BuildDatabase::TargetInfo>> BuildDatabase::getAllTargets() const {
    return CollectionUtils::getValues(targetInfos);
}

std::vector<fs::path> BuildDatabase::getAllTargetPaths() const {
    return CollectionUtils::transformTo<std::vector<fs::path>>(CollectionUtils::getValues(targetInfos),
                                                               [](const std::shared_ptr<TargetInfo> &targetInfo) {
                                                                   return targetInfo->getOutput();
                                                               });
}

std::vector<std::shared_ptr<BuildDatabase::TargetInfo>>
BuildDatabase::getTargetsForSourceFile(const fs::path &sourceFilePath) const {
    CollectionUtils::MapFileTo<bool> cache;
    std::function<
            bool(fs::path const &)> containsSourceFilePath = [&](fs::path const &unitFile) {
        if (CollectionUtils::containsKey(cache, unitFile)) {
            return cache[unitFile];
        }
        if (Paths::isObjectFile(unitFile)) {
            auto compilationUnitInfo = getClientCompilationUnitInfo(unitFile);
            bool isSame = compilationUnitInfo->getSourcePath() == sourceFilePath;
            return cache[unitFile] = isSame;
        }
        auto linkUnitInfo = getClientLinkUnitInfo(unitFile);
        bool result = CollectionUtils::anyTrue(CollectionUtils::transform(
                linkUnitInfo->files, [&containsSourceFilePath](fs::path const &subFile) {
                    return containsSourceFilePath(subFile);
                }));
        return cache[unitFile] = result;
    };

    auto rootTargets = getRootTargets();
    return CollectionUtils::filterOut(
            rootTargets, [&](const std::shared_ptr<const BuildDatabase::TargetInfo> &rootTarget) {
                return !containsSourceFilePath(rootTarget->getOutput());
            });
}

std::vector<fs::path> BuildDatabase::getTargetPathsForSourceFile(const fs::path &sourceFilePath) const {
    auto result = CollectionUtils::transformTo<std::vector<fs::path>>(
            getTargetsForSourceFile(sourceFilePath),
            [&](const std::shared_ptr<const BuildDatabase::TargetInfo> &targetInfo) {
                return targetInfo->getOutput();
            });
    return result;
}

std::vector<fs::path> BuildDatabase::getTargetPathsForObjectFile(const fs::path &objectFile) const {
    std::vector<fs::path> parents;
    if (CollectionUtils::containsKey(objectFileTargets, objectFile)) {
        parents = objectFileTargets.at(objectFile);
    } else {
        LOG_S(WARNING) << "No link unit parents were found for an object file: " << objectFile;
    }
    return parents;
}

std::shared_ptr<BuildDatabase::TargetInfo> BuildDatabase::getPriorityTarget() const {
    CollectionUtils::MapFileTo<int> cache;
    std::function<
            int(fs::path const &)> numberOfSources = [&](fs::path const &unitFile) {
        if (CollectionUtils::containsKey(cache, unitFile)) {
            return cache[unitFile];
        }
        if (Paths::isObjectFile(unitFile)) {
            return 1;
        }
        auto linkUnitInfo = getClientLinkUnitInfo(unitFile);
        int result = 0;
        for (const fs::path &subFile: linkUnitInfo->files) {
            result += numberOfSources(subFile);
        }
        return cache[unitFile] = result;
    };

    auto rootTargets = getRootTargets();
    auto it = std::max_element(rootTargets.begin(), rootTargets.end(),
                               [&](const std::shared_ptr<BuildDatabase::TargetInfo> &a,
                                   const std::shared_ptr<BuildDatabase::TargetInfo> &b) {
                                   return numberOfSources(a->getOutput()) <
                                          numberOfSources(b->getOutput());
                               });
    return *it;
}

fs::path BuildDatabase::newDirForFile(const fs::path &file) const {
    fs::path base = Paths::longestCommonPrefixPath(this->projectContext.getBuildDirAbsPath(),
                                                   this->projectContext.projectPath);
    return Paths::createNewDirForFile(file, base, this->serverBuildDir);
}

CollectionUtils::FileSet BuildDatabase::getSourceFilesForTarget(const fs::path &_target) {
    return CollectionUtils::transformTo<CollectionUtils::FileSet>(
            getArchiveObjectFiles(_target),
            [this](fs::path const &objectPath) {
                return getClientCompilationUnitInfo(objectPath)->getSourcePath();
            });
}

std::shared_ptr<BuildDatabase::TargetInfo> BuildDatabase::getTargetInfo(const fs::path &_target) {
    return targetInfos[_target];
}
